/*
 * ProgramStreamParser.cpp
 *
 *  Created on: 2015年7月5日
 *      Author: terry
 */

#include "ProgramStreamParser.h"

#ifndef WIN32
typedef unsigned char BYTE;
typedef unsigned short WORD;

#define MAKEWORD(a,b) ((WORD) (((BYTE) (a)) | ((WORD) ((BYTE) (b))) << 8))

#endif //WIN32


namespace av
{

static uint8_t START_CODE[] = {0, 0, 1};



ProgramStreamParser::ProgramStreamParser():
m_lastType()
{
    m_cache.ensure(1024 * 128);
    m_pktBuffer.ensure(1024 * 128);
}

ProgramStreamParser::~ProgramStreamParser()
{

}

bool ProgramStreamParser::inputData(const uint8_t* data, size_t length, StreamPacket& pkt)
{
    if (length == 0)
    {
        return false;
    }

    m_cache.write(data, length);

    if (!ensureStartCode(m_cache))
    {
        return false;
    }

    bool found = false;
    while (!found)
    {
        uint8_t* data = m_cache.getReadPtr();
        size_t size = m_cache.readable();

        size_t idx = findHeaderStartCode(data, sizeof(START_CODE), size);
        if (idx == (size_t)-1)
        {
            break;
        }

        StreamPacket packet;
        packet.data = data;
        packet.length = idx;
        packet.type = data[3];

        found = handlePacket(packet);

        m_cache.skip(idx);
    }

    if (found)
    {
        pkt = m_naluPkt;
    }

    return found;
}


void ProgramStreamParser::clear()
{
    m_cache.clear();
}


bool ProgramStreamParser::handlePacket(StreamPacket& packet)
{
	bool found = false;
    if (packet.type >= 0xBA && packet.type <= 0xBF)    // PS header
    {
    	found = flushPacketBuffer();
    }
    else if (packet.type >= 0xC0 && packet.type <= 0xDF) // audio
    {
        /// 忽略音频包
    	found = false;
    }
    else if (packet.type >= 0xE0 && packet.type <= 0xEF)    // video stream
    {
        found = onPESPacket(packet);
    }
    else
    {
        /// 忽略不识别的包
    	found = false;
    }
    return found;
}


bool ProgramStreamParser::onPESPacket(StreamPacket& packet)
{
    if (packet.length < (sizeof(START_CODE) + 3))
    {
        return false;
    }

    uint16_t pktLength = MAKEWORD(packet.data[5], packet.data[4]);
    size_t hdrLength = packet.data[5 + 3];
    size_t offset = sizeof(START_CODE) + 3 + 3 + hdrLength;

    if (packet.length < offset)
    {
        return false;
    }

    bool found = false;
    uint8_t* data = packet.data + offset;
	uint16_t length = (uint16_t)(packet.length - offset);

	if (H264NaluParser::startWithH264Code(data, length))
	{
	    /// 新的NALU包
	    found = flushPacketBuffer();

	    m_pktBuffer.write(data, length);
	}
	else
	{
	    m_pktBuffer.write(data, length);
	}
	return found;
}

void ProgramStreamParser::writePacket(NaluPacket& packet)
{
	m_naluBuffer.clear();

	m_naluBuffer.write(packet.data - packet.prefix, packet.length + packet.prefix);

	m_naluPkt.data = m_naluBuffer.getReadPtr();
	m_naluPkt.length = m_naluBuffer.readable();
	m_naluPkt.type = packet.type;
}

bool ProgramStreamParser::flushPacketBuffer()
{
    if (m_pktBuffer.empty())
    {
        return false;
    }

    NaluPacket pkt;
    if (!H264NaluParser::parseNalu(m_pktBuffer.getReadPtr(), m_pktBuffer.readable(), pkt))
    {
        return false;
    }

    writePacket(pkt);

    m_pktBuffer.clear();

    return true;
}

size_t ProgramStreamParser::findStartCode(const uint8_t* buffer, size_t length)
{
    if (length < sizeof(START_CODE))
    {
        return -1;
    }

    for (size_t i = 0; i < (length - sizeof(START_CODE)); ++ i)
    {
        if ((buffer[i] == 0) && (buffer[i+1] == 0) && (buffer[i+2] == 1))
        {
            return i;
        }
    }
    return -1;
}

size_t ProgramStreamParser::findStartCode(const uint8_t* buffer, size_t offset, size_t length)
{
    if (length <= offset)
    {
        return -1;
    }

    size_t idx = findStartCode(buffer + offset, length - offset);
    if (idx != (size_t)-1)
    {
        idx += offset;
    }
    return idx;
}

size_t ProgramStreamParser::findHeaderStartCode(const uint8_t* buffer, size_t offset, size_t length)
{
    size_t idx = findStartCode(buffer, offset, length);
    while (idx != (size_t)-1)
    {
        uint8_t type = buffer[idx + sizeof(START_CODE)];
        if (type >= 0xBA)
        {
            break;
        }
        else
        {
            idx = findStartCode(buffer, idx + sizeof(START_CODE), length);
        }
    }
    return idx;
}

bool ProgramStreamParser::ensureStartCode(comn::IOBuffer& buffer)
{
    bool found = false;
    uint8_t* data = m_cache.getReadPtr();
    size_t size = m_cache.readable();
    size_t idx = findStartCode(data, size);
    if (idx == (size_t)-1)
    {
        if (size > sizeof(START_CODE))
        {
            buffer.skip(size - sizeof(START_CODE));
        }
    }
    else
    {
        buffer.skip(idx);
        found = true;
    }
    return found;
}




} /* namespace av */
